/* Copyright (C) 2021 Rain */

/* This file is part of XNIX. */

/* 
  XNIX is free software: you can redistribute it and/or modify 
  it under the terms of the GNU General Public License as published by 
  the Free Software Foundation, either version 3 of the License, or 
  (at your option) and later version. 
*/

/*
  XNIX is distributed in the hope that it will be useful, 
  but WITHOUT ANY WARRANTY; without even the implied warranty of 
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 
  GNU General Public License for more details. 
*/

/*
  You should have received a copy of the GNU General Public License 
   along with XNIX. If not, see <https://www.gnu.org/licenses/>.  
*/




#include <kernel/init.h>
#include <kernel/errno.h>
#include <kernel/types.h>

#include <kernel/console.h>
#include <kernel/int.h>
#include <kernel/panic.h>
#include <kernel/mm.h>

#include <lib/bitmap.h>
#include <lib/string.h>

#include <arch/x86/int.h>
#include <arch/x86/paging.h>

#include <arch/io.h>
#include <arch/sys.h>


struct sysmm *mm;

struct mempool pool, upool, devpool;


void mm_init(void) {
	mm = (struct sysmm *) PAGE_4K_ALIGN((__uint64_t) &_end);
	
	struct e820_entry *e820 = E820_ADDR;
	struct zone *zones = (struct zone *) (((char *) mm) + sizeof(struct sysmm));

	int i;
	__uint64_t total_ram = 0;

	struct e820_entry *e = e820;

	memset(mm, 0x00, sizeof(struct sysmm));

	for (i = 0; i < 32; i++) {
		if (!e->type || e->type >= E820_INV) 
			break;

		printk(white, "address %q, length %q KB, type %d\n", e->addr, e->len >> 10, e->type);

		/* create a `struct zone` for this */

#define z (zones + mm->total_zones)

		z->type = e->type;
		z->total_pages = e->len >> PAGE_SHIFT;

		z->start_addr = e->addr;
		z->end_addr = e->addr + e->len;

		z->size = e->len;
		z->using = 0;

		mm->total_zones++;

		if (e->type == E820_RAM) 
			total_ram += e->len;

		mm->total_pages += e->len >> PAGE_SHIFT;
		e++;
	}

	printk(white, "total usable ram: %d MB. \n", total_ram >> 20);
	printk(white, "total pages: %d (includes unusable pages) = %d MB\n", mm->total_pages, mm->total_pages << PAGE_SHIFT >> 20);

	printk(white, "total zones: %d\n", mm->total_zones);

	/* setup page bitmap */

	/* first we should allocate some memory for it */
	/* at end of zones */

	mm->bitmap = (__uint8_t *) (zones + mm->total_zones);

	/* set all bits. if this bit points to a RAM page, clear it */

	/* it always not bad. */
#define BYTESNEED ((mm->total_pages >> 3) + 1)

	memset(mm->bitmap, 0xff, BYTESNEED);

	e = e820;

	for (i = 0; i < mm->total_zones; i++) {

#undef z
#define z (zones + i)

		if (z->type == E820_RAM) {
			/* clear bits in bitmap */
			int j;
			__uint64_t start_page = PAGE_ALIGN(z->start_addr) >> PAGE_SHIFT;

			for (j = start_page; j < start_page + z->total_pages; j++) 
				bitmap_set(mm->bitmap, j, 0);
		}
	}

	/* 
	 * the first and second page is for kernel, 
	 * and the 3rd, 4th, 5th are for frame buffer (they setup in boot/loader.asm). 
	 * */

	/* 0MB */
	bitmap_set(mm->bitmap, 0, 1);

	/* 2MB */
	bitmap_set(mm->bitmap, 1, 1);

	/* now we setup that memory pool */
	/* default pool size: 2MB (1 page) */

	pool.size = 0x200000;
	pool.used = 0;

	pool.entries = 0x200000 >> POOL_ENTRY_SHIFT;

	/* 32 bytes per bit */

	/* we don't want use alloc_page() here, because it just want 
	 * 4 KB, but alloc_page() gives us 2MB. 
	 *
	 * so we must find a place to store this. 
	 *
	 * and after that we can use kmalloc() to alloc this. 
	 * (this is only for starting-up)
	 *
	 * it's at the end of syslog-buffer. 
	 * */

	pool.bitmap = (__uint8_t *) (SYSLOG_BUFFER + 32768);

	memset(pool.bitmap, 0, 0x200000 / 64 / 8);

	pool.page = alloc_page();

	if (!pool.page) 
		panic("mm_init: fatal allocate page for pool. \n");

	upool.size = 0x200000;
	upool.used = 0;

	upool.entries = 0x200000 >> POOL_ENTRY_SHIFT;

	upool.bitmap = (__uint8_t *) kmalloc(&pool, 4096);

	if (!upool.bitmap) 
		panic("mm_init: fatal allocate space for upool. ");

	memset(upool.bitmap, 0, 0x200000 / 64 / 8);

	upool.page = alloc_page_att(PATT_USER);

	if (!upool.page) 
		panic("mm_init: fatal allocate page for upool. ");

	pool.next = &upool;
	
	devpool.size = 0x200000;
	devpool.used = 0;

	devpool.entries = 0x200000 >> POOL_ENTRY_SHIFT;

	devpool.bitmap = (__uint8_t *) kmalloc(&pool, 4096);

	if (!devpool.bitmap) 
		panic("mm_init: fatal allocate space for devpool. ");

	memset(devpool.bitmap, 0, 0x200000 / 64 / 8);

	devpool.page = alloc_page_att(PATT_DEVMM);

	if (!devpool.page) 
		panic("mm_init: fatal allocate page for devpool. ");

	upool.next = &devpool;

	return;

#undef z
}

/* find an empty page table entry, and return its number */

void *mapping(void *p, __uint64_t attr) {
	__uint64_t i;
	void *addr = (void *) ((__uint64_t) p & PAGE_MASK);

	/* we just use the (max) 1 GB memory in kernel (512 pages) */

#define PT_ADDR	((__uint64_t *) 0x2000)

	for (i = 0; i < 512; i++) 
		if (PT_ADDR[i] == (__uint64_t) addr) 
			return (void *) ((i << PAGE_SHIFT) + p - addr);

	for (i = 0; i < 512; i++) {
		if (!PT_ADDR[i]) {
			PT_ADDR[i] = (__uint64_t) addr | attr;
			invalidate();

			return (void *) ((i << PAGE_SHIFT) + p - addr);
		}
	}

	return NULL;
#undef PT_ADDR
}


__uint64_t unmapping(void *p) {
	int i;
	__uint64_t addr = (__uint64_t) p & PAGE_MASK;

#define PT_ADDR	((__uint64_t *) 0x2000)

	for (i = 0; i < 512; i++) 
		if (PT_ADDR[i] == addr) 
			return PT_ADDR[i] & 0xff;

	panic("unmapping: page not found. ");

	return 0;

}


void *alloc_page(void) {
	return alloc_page_att(PATT_KERNEL);
}


void *alloc_page_att(__uint64_t attr) {
	__uint64_t i;

	for (i = 0; bitmap_read(mm->bitmap, i); i++) {
		if (i > mm->total_pages) 
			return NULL;
	}

	bitmap_set(mm->bitmap, i, 1);

	/* we must map it to page-table */

	void *addr = mapping((void *) (i << PAGE_SHIFT), attr);

	if (!addr) {
		bitmap_set(mm->bitmap, i, 0);
		return NULL;
	}

	return addr;
}
 

void free_page(void *p) {
	__uint64_t i;

	i = (__uint64_t) p >> PAGE_SHIFT;

	bitmap_set(mm->bitmap, i, 0);

	return;
}


void *kmalloc(struct mempool *pool, __uint64_t size) {
	if (!size || size > pool->size - pool->used) 
		return NULL;

	/* devpool needs kb-aligning */

	if (pool == &devpool) 
		size = (size + 1023) & ~1023;

	__uint64_t points = ((size + 63) & ~63) >> POOL_ENTRY_SHIFT;
	__uint64_t i;

	if (points == 1) {
		for (i = 0; bitmap_read(pool->bitmap, i); i++) {
			if (i > pool->entries) 
				goto panic;
		}

		bitmap_set(pool->bitmap, i, 1);
		pool->used++;

		return pool->page + (i << POOL_ENTRY_SHIFT);
	}

	/* it's not so good... */

	for (i = 0; i < pool->entries; i++) {

repeat:
		/* if all 0... */
		__uint64_t j;

		for (j = 0; j < points; j++) {
			if (bitmap_read(pool->bitmap, i + j)) {
				i++;

				if (i == pool->entries) 
					goto panic;

				goto repeat;
			}
		}

		/* it will here when we found enough */

		for (j = 0; j < points; j++) 
			bitmap_set(pool->bitmap, i + j, 1);

		pool->used += points;

		return pool->page + (i << POOL_ENTRY_SHIFT);
	}

panic:

	printk(yellow, "kmalloc: no enough memory (total %d KB, %d bytes used, requested %d bytes)\n", 
		pool->size >> 10, pool->used << 6, points << 6);

	return NULL;
}


void kfree(struct mempool *pool, void *p, __uint64_t size) {
	if (pool == &devpool) 
		size = (size + 1023) & ~1023;
	
	__uint64_t points = ((size + 63) & ~63) >> POOL_ENTRY_SHIFT;
	__uint64_t i;
	
	i = (p - pool->page) >> POOL_ENTRY_SHIFT;

	for (; points; points--, i++) 
		bitmap_set(pool->bitmap, i, 0);

	pool->used -= points;

	return;
}


void trace_mm(void) {
	printk(yellow, "tracing system memory: \n");

	__uint64_t *cr3 = (__uint64_t *) get_cr3();

	printk(yellow, "pml4: 0x%q  pdp: 0x%q  pt: 0x%q\n", 
		cr3, *cr3 & ~0xff, *(__uint64_t *) (*cr3 & ~0xff) & ~0xff);

	__uint64_t *pt = (__uint64_t *) (*(__uint64_t *) (*cr3 & ~0xff) & ~0xff);
	int i;

	for (i = 0; i < 512; i++) {
		if (pt[i] & ~0xff) 
			printk(yellow, "\t%q -> %q: %b\n", i << PAGE_SHIFT, 
					pt[i] & ~0xff, pt[i] & 0xff);
	}
		
	return;
}


/* 
 * get the physical address of virtual address 'p'
 * the `mapping()` got the vaddress of physical, and we will 
 * get the physical (find in the page-table). 
 * */

void *mm_vphy(void *p) {
	/* like mapping(), we just search in the first 1g */

	__uint64_t addr = (__uint64_t) p & PAGE_MASK;

#define PT_ADDR ((__uint64_t *) 0x2000)

	__uint64_t i;

	for (i = 0; i < 512; i++) 
		if ((PT_ADDR[i] & ~0xff) == addr) 
			return (void *) ((i << PAGE_SHIFT) + ((__uint64_t) p - addr));

	return NULL;

#undef PT_ADDR
}

